ZeRO: Zero Redundancy Optimizer
ZeRO adopts different levels of distributed strategy to remove redundancy in memory, at the cost of more communication. This work also gives birth to FSDP, a widely used distributed training schema.
Table of Contents
1. GPU Memory Breakdown
When training a model on GPU, the GPU memory can be divided into 6 parts.
- Model parameters
- Gradients
- Optimizer states
- Activations, e.g., the mean and std required for
Layernorm’s backprop - Temporary buffers, workspace
- Fragmentation
For example, training a 7B model in FP16 (2 bytes).
- Memory for parameters is \(7B \times 2\) bytes
- Gradients of parameters are \(7B \times 2\) bytes
- Suppose we use Adam optimizer, for each weight, we need to store a FP32 copy of the weight (for numerical stability), and \(m,v\) in FP32. That is \(7B \times (4+4+4)\) bytes
Even if we have already ignored activations and all other consumptions, training a 7B model already consumes \(112\) GB memory.
2. Methodology
In classic distributed data parallelism (DDP), each GPU has a full copy of the model and processes a fraction of data. After each step of backprop, gradients are allreduce-ed for updating the parameters. Each rank in DDP have a full and identical copy of parameters, gradients and optimizer states.
ZeRO handles the redundancy problem in DDP by eliminating redundant copies of parameters, gradients, and optimizer states at the cost of communication.
ZeRO can be divided into 3 levels, based on aggressiveness of elimination.
- ZeRO-1. Only cut optimizer states.
- ZeRO-2. Optimizer states + gradients.
- ZeRO-3. Optimizer states + gradients + parameters.
It’s worth noticing that FSDP is PyTorch’s implementation of ZeRO-3.
2.1. ZeRO-1: Optimizer States Sharding
ZeRO-1 shards optimizer states by DP size. On each rank, the parameters and gradients are still full copy, but optimizer states are sharded.